

"""
用于可视化的一组函数。
这些功能通常接收图像，对图像进行一些可视化处理。
这些函数不返回值，而是修改图像本身。
"""
import collections
import numpy as np
import PIL.Image as Image
import PIL.ImageColor as ImageColor
import PIL.ImageDraw as ImageDraw
import PIL.ImageFont as ImageFont
import six
import tensorflow as tf


_TITLE_LEFT_MARGIN = 10
_TITLE_TOP_MARGIN = 10
STANDARD_COLORS = [
    'AliceBlue', 'Chartreuse', 'Aqua', 'Aquamarine', 'Azure', 'Beige', 'Bisque',
    'BlanchedAlmond', 'BlueViolet', 'BurlyWood', 'CadetBlue', 'AntiqueWhite',
    'Chocolate', 'Coral', 'CornflowerBlue', 'Cornsilk', 'Crimson', 'Cyan',
    'DarkCyan', 'DarkGoldenRod', 'DarkGrey', 'DarkKhaki', 'DarkOrange',
    'DarkOrchid', 'DarkSalmon', 'DarkSeaGreen', 'DarkTurquoise', 'DarkViolet',
    'DeepPink', 'DeepSkyBlue', 'DodgerBlue', 'FireBrick', 'FloralWhite',
    'ForestGreen', 'Fuchsia', 'Gainsboro', 'GhostWhite', 'Gold', 'GoldenRod',
    'Salmon', 'Tan', 'HoneyDew', 'HotPink', 'IndianRed', 'Ivory', 'Khaki',
    'Lavender', 'LavenderBlush', 'LawnGreen', 'LemonChiffon', 'LightBlue',
    'LightCoral', 'LightCyan', 'LightGoldenRodYellow', 'LightGray', 'LightGrey',
    'LightGreen', 'LightPink', 'LightSalmon', 'LightSeaGreen', 'LightSkyBlue',
    'LightSlateGray', 'LightSlateGrey', 'LightSteelBlue', 'LightYellow', 'Lime',
    'LimeGreen', 'Linen', 'Magenta', 'MediumAquaMarine', 'MediumOrchid',
    'MediumPurple', 'MediumSeaGreen', 'MediumSlateBlue', 'MediumSpringGreen',
    'MediumTurquoise', 'MediumVioletRed', 'MintCream', 'MistyRose', 'Moccasin',
    'NavajoWhite', 'OldLace', 'Olive', 'OliveDrab', 'Orange', 'OrangeRed',
    'Orchid', 'PaleGoldenRod', 'PaleGreen', 'PaleTurquoise', 'PaleVioletRed',
    'PapayaWhip', 'PeachPuff', 'Peru', 'Pink', 'Plum', 'PowderBlue', 'Purple',
    'Red', 'RosyBrown', 'RoyalBlue', 'SaddleBrown', 'Green', 'SandyBrown',
    'SeaGreen', 'SeaShell', 'Sienna', 'Silver', 'SkyBlue', 'SlateBlue',
    'SlateGray', 'SlateGrey', 'Snow', 'SpringGreen', 'SteelBlue', 'GreenYellow',
    'Teal', 'Thistle', 'Tomato', 'Turquoise', 'Violet', 'Wheat', 'White',
    'WhiteSmoke', 'Yellow', 'YellowGreen'
]

#将表示图像的 numpy 数组保存为 PNG 格式的图像文件。
def save_image_array_as_png(image, output_path):
  """
  将图像（表示为numpy数组）保存为PNG。
  Args:
    image：形状为[高，宽，3]的numpy数组。
    output_path：应该写入图像的路径。
  """
  image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
  with tf.gfile.Open(output_path, 'w') as fid:
    image_pil.save(fid, 'PNG')

#将 numpy 数组编码为 PNG 格式的字符串。
def encode_image_array_as_png_str(image):
  """
    将numpy数组编码为PNG字符串。
  Args:
    image：形状为[高，宽，3]的numpy数组。
  Returns:
    PNG编码的图像字符串。
  """
  image_pil = Image.fromarray(np.uint8(image))
  output = six.BytesIO()
  image_pil.save(output, format='PNG')
  png_string = output.getvalue()
  output.close()
  return png_string

# 在图像的 numpy 数组上添加边界框。可以指定边界框的位置、颜色、厚度以及要在边界框内显示的文本
def draw_bounding_box_on_image_array(image,
                                     ymin,
                                     xmin,
                                     ymax,
                                     xmax,
                                     color='red',
                                     thickness=4,
                                     display_str_list=(),
                                     use_normalized_coordinates=True):
  """将边界框添加到图像（numpy数组）。
  Args:
        image：形状为[高，宽，3]的numpy数组。
        ymin:ymin在归一化坐标中的边界框（下同）。
        xmin：边界框的xmin。
        ymax:ymax的边界框。
        xmax:xmax的边界框。
        color：绘制边框的颜色。默认值为红色。
        thickness：线条厚度。默认值为4。
        display_str_list：要在框中显示的字符串列表（每一个都显示在自己的行上）。
        use_normalized_cordinates：如果为True（默认值），则将坐标 ymin、xmin、ymax、xmax 视为相对于图像的比例坐标。否则将坐标视为绝对坐标。
  """

  image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
  draw_bounding_box_on_image(image_pil, ymin, xmin, ymax, xmax, color,
                             thickness, display_str_list,
                             use_normalized_coordinates)
  np.copyto(image, np.array(image_pil))

  """draw_bounding_box_on_image: 将边界框添加到图像中。每个字符串在 display_str_list 中都会以黑色文本显示在边界框上方的单独行中，
    这些行位于一个填充了指定颜色的矩形内。
            image: PIL.Image 对象，表示要绘制边界框的图像。
            ymin: 边界框的 ymin 坐标。
            xmin: 边界框的 xmin 坐标。
            ymax: 边界框的 ymax 坐标。
            xmax: 边界框的 xmax 坐标。
            color: 绘制边界框的颜色。默认为红色。
            thickness: 边界框线条的粗细。默认值为 4。
            display_str_list: 要显示在边界框内的字符串列表，每个字符串将显示在独立的行上。默认为空列表。
            use_normalized_coordinates: 如果为 True（默认），则将坐标 ymin、xmin、ymax、xmax 视为相对于图像的比例坐标。否则将坐标视为绝对坐标。"""
def draw_bounding_box_on_image(image,
                               ymin,
                               xmin,
                               ymax,
                               xmax,
                               color='red',
                               thickness=4,
                               display_str_list=(),
                               use_normalized_coordinates=True):
  draw = ImageDraw.Draw(image)
  im_width, im_height = image.size
  if use_normalized_coordinates:
    (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                  ymin * im_height, ymax * im_height)
  else:
    (left, right, top, bottom) = (xmin, xmax, ymin, ymax)
  draw.line([(left, top), (left, bottom), (right, bottom),
             (right, top), (left, top)], width=thickness, fill=color)
  try:
    font = ImageFont.truetype('arial.ttf', 24)
  except IOError:
    font = ImageFont.load_default()

  text_bottom = top
  # Reverse list and print from bottom to top.
  for display_str in display_str_list[::-1]:
    text_width, text_height = font.getsize(display_str)
    margin = np.ceil(0.05 * text_height)
    draw.rectangle(
        [(left, text_bottom - text_height - 2 * margin), (left + text_width,
                                                          text_bottom)],
        fill=color)
    draw.text(
        (left + margin, text_bottom - text_height - margin),
        display_str,
        fill='black',
        font=font)
    text_bottom -= text_height - 2 * margin

"""在图像的 numpy 数组上绘制多个边界框。
image: 一个 numpy 数组，表示要绘制边界框的图像。
boxes: 边界框的列表，每个边界框由四个坐标（ymin、xmin、ymax、xmax）组成。坐标的规格化格式介于[0，1]之间。
color: 绘制边界框的颜色。默认为红色。
thickness: 边界框线条的粗细。默认值为 4。
display_str_list_list: 要显示在每个边界框内的字符串列表的列表。默认为空列表。"""
def draw_bounding_boxes_on_image_array(image,
                                       boxes,
                                       color='red',
                                       thickness=4,
                                       display_str_list_list=()):

  image_pil = Image.fromarray(image)
  draw_bounding_boxes_on_image(image_pil, boxes, color, thickness,
                               display_str_list_list)
  np.copyto(image, np.array(image_pil))

'''在图像上绘制边界框。
image: 一个 PIL.Image 对象。
boxes: 一个二维 numpy 数组，形状为 [N, 4]，表示 N 个边界框的坐标 (ymin, xmin, ymax, xmax)，坐标值在 [0, 1] 范围内。
color: 绘制边界框的颜色。默认为红色。
thickness: 边界框线条的粗细。默认值为 4。
display_str_list_list: 字符串列表的列表，每个边界框可能包含多个标'''
def draw_bounding_boxes_on_image(image,
                                 boxes,
                                 color='red',
                                 thickness=4,
                                 display_str_list_list=()):

  boxes_shape = boxes.shape
  if not boxes_shape:
    return
  if len(boxes_shape) != 2 or boxes_shape[1] != 4:
    raise ValueError('Input must be of size [N, 4]')
  for i in range(boxes_shape[0]):
    display_str_list = ()
    if display_str_list_list:
      display_str_list = display_str_list_list[i]
    draw_bounding_box_on_image(image, boxes[i, 0], boxes[i, 1], boxes[i, 2],
                               boxes[i, 3], color, thickness, display_str_list)

'''在图像的 numpy 数组上绘制关键点。
image: 一个 numpy 数组，表示要绘制关键点的图像。
keypoints: 一个 numpy 数组，形状为 [num_keypoints, 2]，表示关键点的坐标。
color: 绘制关键点的颜色。默认为红色。
radius: 关键点的半径。默认值为 2。
use_normalized_coordinates: 如果为 True（默认），则将关键点的值视为相对于图像的比例坐标。否则将坐标视为绝对坐标。'''
def draw_keypoints_on_image_array(image,
                                  keypoints,
                                  color='red',
                                  radius=2,
                                  use_normalized_coordinates=True):

  image_pil = Image.fromarray(np.uint8(image)).convert('RGB')
  draw_keypoints_on_image(image_pil, keypoints, color, radius,
                          use_normalized_coordinates)
  np.copyto(image, np.array(image_pil))

'''在 PIL.Image 对象上绘制关键点。
参数与 draw_keypoints_on_image_array 相同，只是 image 的类型为 PIL.Image 对象。'''
"""在图像上绘制边界框和关键点，实现对检测和识别结果的可视化显示。"""
def draw_keypoints_on_image(image,
                            keypoints,
                            color='red',
                            radius=2,
                            use_normalized_coordinates=True):
  """
      在图像上绘制关键点。
        Args：
        图片： PIL。图像对象。
        keypoints：形状为[num_keypoints，2]的numpy数组。
        color：绘制关键点的颜色。默认值为红色。
        radius：关键点半径。默认值为2。
        use_normalized_cordinates：如果为True（默认值），则将关键点值视为相对于图像。否则，就把它们当作绝对的。
  """
  draw = ImageDraw.Draw(image)
  im_width, im_height = image.size
  keypoints_x = [k[1] for k in keypoints]
  keypoints_y = [k[0] for k in keypoints]
  if use_normalized_coordinates:
    keypoints_x = tuple([im_width * x for x in keypoints_x])
    keypoints_y = tuple([im_height * y for y in keypoints_y])
  for keypoint_x, keypoint_y in zip(keypoints_x, keypoints_y):
    draw.ellipse([(keypoint_x - radius, keypoint_y - radius),
                  (keypoint_x + radius, keypoint_y + radius)],
                 outline=color, fill=color)


'''在图像的 numpy 数组上绘制掩码。
image: uint8 类型的 numpy 数组，形状为 (img_height, img_height, 3)，表示图像。
mask: 形状为 (img_height, img_height) 的浮点型 numpy 数组，值在 [0, 1] 范围内。
color: 绘制掩码的颜色。默认为红色。
alpha: 透明度值，范围在 0 到 1 之间。默认值为 0.7。
'''
def draw_mask_on_image_array(image, mask, color='red', alpha=0.7):
  """
    Raises：
    ValueError：图像或掩码的数据类型不正确。
  """
  if image.dtype != np.uint8:
    raise ValueError('`image` not of type np.uint8')
  if mask.dtype != np.float32:
    raise ValueError('`mask` not of type np.float32')
  if np.any(np.logical_or(mask > 1.0, mask < 0.0)):
    raise ValueError('`mask` elements should be in [0, 1]')
  rgb = ImageColor.getrgb(color)
  pil_image = Image.fromarray(image)

  solid_color = np.expand_dims(
      np.ones_like(mask), axis=2) * np.reshape(list(rgb), [1, 1, 3])
  pil_solid_color = Image.fromarray(np.uint8(solid_color)).convert('RGBA')
  pil_mask = Image.fromarray(np.uint8(255.0*alpha*mask)).convert('L')
  pil_image = Image.composite(pil_solid_color, pil_image, pil_mask)
  np.copyto(image, np.array(pil_image.convert('RGB')))


"""在图像的 numpy 数组上叠加带有格式化得分和标签名称的标注框。
image: uint8 类型的 numpy 数组，形状为 (img_height, img_width, 3)，表示图像。
boxes: 形状为 [N, 4] 的 numpy 数组，表示 N 个边界框的坐标 (ymin, xmin, ymax, xmax)。
classes: 形状为 [N] 的 numpy 数组，表示 N 个边界框的类别。
scores: 形状为 [N] 的 numpy 数组或 None。如果 scores=None，则假定要绘制的边界框是 groundtruth 边界框，并且所有边界框都会绘制为黑色，没有类别或得分。
category_index: 包含类别字典（每个字典包含类别索引 id 和类别名称 name）的字典，以类别索引为键。
instance_masks: 形状为 [N, image_height, image_width] 的 numpy 数组，可以为 None。
keypoints: 形状为 [N, num_keypoints, 2] 的 numpy 数组，可以为 None。
use_normalized_coordinates: 是否将边界框解释为标准化坐标。默认为 False。
max_boxes_to_draw: 最大要可视化的边界框数。如果为 None，则绘制所有边界框。
min_score_thresh: 要可视化的边界框的最小得分阈值。
agnostic_mode: 布尔值，默认为 False，控制是否以类别不可知模式评估。该模式将显示得分但忽略类别。
line_thickness: 整数，默认为 4，控制边界框的线宽度"""
def visualize_boxes_and_labels_on_image_array(image,
                                              boxes,
                                              classes,
                                              scores,
                                              category_index,
                                              instance_masks=None,
                                              keypoints=None,
                                              use_normalized_coordinates=False,
                                              max_boxes_to_draw=20,
                                              min_score_thresh=.7,
                                              agnostic_mode=False,
                                              line_thickness=4):

  # 为每个框位置创建一个显示字符串（和颜色），对任何框进行分组
  # 它们对应于相同的位置。
  box_to_display_str_map = collections.defaultdict(list)
  box_to_color_map = collections.defaultdict(str)
  box_to_instance_masks_map = {}
  box_to_keypoints_map = collections.defaultdict(list)
  if not max_boxes_to_draw:
    max_boxes_to_draw = boxes.shape[0]
  for i in range(min(max_boxes_to_draw, boxes.shape[0])):
    if scores is None or scores[i] > min_score_thresh:
      box = tuple(boxes[i].tolist())
      if instance_masks is not None:
        box_to_instance_masks_map[box] = instance_masks[i]
      if keypoints is not None:
        box_to_keypoints_map[box].extend(keypoints[i])
      if scores is None:
        box_to_color_map[box] = 'black'
      else:
        if not agnostic_mode:
          if classes[i] in category_index.keys():
            class_name = category_index[classes[i]]['name']
          else:
            class_name = 'N/A'
          display_str = '{}: {}%'.format(
              class_name,
              int(100*scores[i]))
        else:
          display_str = 'score: {}%'.format(int(100 * scores[i]))
        box_to_display_str_map[box].append(display_str)
        if agnostic_mode:
          box_to_color_map[box] = 'DarkOrange'
        else:
          box_to_color_map[box] = STANDARD_COLORS[
              classes[i] % len(STANDARD_COLORS)]

  #将所有框绘制到图像上
  for box, color in box_to_color_map.items():
    color = 'Violet'
    ymin, xmin, ymax, xmax = box
    if instance_masks is not None:
      draw_mask_on_image_array(
          image,
          box_to_instance_masks_map[box],
          color=color
      )
    draw_bounding_box_on_image_array(
        image,
        ymin,
        xmin,
        ymax,
        xmax,
        color=color,
        thickness=line_thickness,
        display_str_list=box_to_display_str_map[box],
        use_normalized_coordinates=use_normalized_coordinates)
    if keypoints is not None:
      draw_keypoints_on_image_array(
          image,
          box_to_keypoints_map[box],
          color=color,
          radius=line_thickness / 2,
          use_normalized_coordinates=use_normalized_coordinates)